import GUI from 'lil-gui';
import * as THREE from 'three';


const waterFrag= /*glsl*/`precision highp float;

uniform float uOpacity;

uniform vec3 uTroughColor;
uniform vec3 uSurfaceColor;
uniform vec3 uPeakColor;

uniform float uPeakThreshold;
uniform float uPeakTransition;
uniform float uTroughThreshold;
uniform float uTroughTransition;

uniform float uFresnelScale;
uniform float uFresnelPower;

varying vec3 vNormal;
varying vec3 vWorldPosition;

uniform samplerCube uEnvironmentMap;

void main() {
  // Calculate vector from camera to the vertex
  vec3 viewDirection = normalize(vWorldPosition - cameraPosition);
  vec3 reflectedDirection = reflect(viewDirection, vNormal);
  reflectedDirection.x = -reflectedDirection.x;

  // Sample environment map to get the reflected color
  vec4 reflectionColor = textureCube(uEnvironmentMap, reflectedDirection);

  // Calculate fresnel effect
  float fresnel = uFresnelScale * pow(1.0 - clamp(dot(viewDirection, vNormal), 0.0, 1.0), uFresnelPower);

  // Calculate elevation-based color
  float elevation = vWorldPosition.y;

  // Calculate transition factors using smoothstep
  float peakFactor = smoothstep(uPeakThreshold - uPeakTransition, uPeakThreshold + uPeakTransition, elevation);
  float troughFactor = smoothstep(uTroughThreshold - uTroughTransition, uTroughThreshold + uTroughTransition, elevation);

  // Mix between trough and surface colors based on trough transition
  vec3 mixedColor1 = mix(uTroughColor, uSurfaceColor, troughFactor);

  // Mix between surface and peak colors based on peak transition 
  vec3 mixedColor2 = mix(mixedColor1, uPeakColor, peakFactor);

  // Mix the final color with the reflection color
  vec3 finalColor = mix(mixedColor2, reflectionColor.rgb, fresnel);

  gl_FragColor = vec4(finalColor, uOpacity);
}
`;
const waterVert = /*glsl*/`precision highp float;

uniform float uTime;

uniform float uWavesAmplitude;
uniform vec2 uWavesSpeed;
uniform float uWavesFrequency;
uniform float uWavesPersistence;
uniform float uWavesLacunarity;
uniform float uWavesIterations;

varying vec3 vNormal;
varying vec3 vWorldPosition;

//	Simplex 3D Noise 
//	by Ian McEwan, Stefan Gustavson (https://github.com/stegu/webgl-noise)
//
vec4 permute(vec4 x) {
  return mod(((x * 34.0) + 1.0) * x, 289.0);
}
vec4 taylorInvSqrt(vec4 r) {
  return 1.79284291400159 - 0.85373472095314 * r;
}

// Simplex 2D noise
//
vec3 permute(vec3 x) {
  return mod(((x * 34.0) + 1.0) * x, 289.0);
}

float snoise(vec2 v) {
  const vec4 C = vec4(0.211324865405187, 0.366025403784439, -0.577350269189626, 0.024390243902439);
  vec2 i = floor(v + dot(v, C.yy));
  vec2 x0 = v - i + dot(i, C.xx);
  vec2 i1;
  i1 = (x0.x > x0.y) ? vec2(1.0, 0.0) : vec2(0.0, 1.0);
  vec4 x12 = x0.xyxy + C.xxzz;
  x12.xy -= i1;
  i = mod(i, 289.0);
  vec3 p = permute(permute(i.y + vec3(0.0, i1.y, 1.0)) + i.x + vec3(0.0, i1.x, 1.0));
  vec3 m = max(0.5 - vec3(dot(x0, x0), dot(x12.xy, x12.xy), dot(x12.zw, x12.zw)), 0.0);
  m = m * m;
  m = m * m;
  vec3 x = 2.0 * fract(p * C.www) - 1.0;
  vec3 h = abs(x) - 0.5;
  vec3 ox = floor(x + 0.5);
  vec3 a0 = x - ox;
  m *= 1.79284291400159 - 0.85373472095314 * (a0 * a0 + h * h);
  vec3 g;
  g.x = a0.x * x0.x + h.x * x0.y;
  g.yz = a0.yz * x12.xz + h.yz * x12.yw;
  return 130.0 * dot(m, g);
}

// Helper function to calculate elevation at any point
float getElevation(float x, float z) {
  vec2 pos = vec2(x, z);

  float elevation = 0.0;
  float amplitude = 1.0;
  float frequency = uWavesFrequency;
  vec2 p = pos.xy;

  for(float i = 0.0; i < uWavesIterations; i++) {
    float noiseValue = snoise(p * frequency + uTime * uWavesSpeed);
    elevation += amplitude * noiseValue;
    amplitude *= uWavesPersistence;
    frequency *= uWavesLacunarity;
  }

  elevation *= uWavesAmplitude;

  return elevation;
}

void main() {
  vec4 modelPosition = modelMatrix * vec4(position, 1.0);

  float elevation = getElevation(modelPosition.x, modelPosition.z);
  modelPosition.y += elevation;

  // Calculate normal using partial derivatives
  float eps = 0.001;
  vec3 tangent = normalize(vec3(eps, getElevation(modelPosition.x - eps, modelPosition.z) - elevation, 0.0));
  vec3 bitangent = normalize(vec3(0.0, getElevation(modelPosition.x, modelPosition.z - eps) - elevation, eps));
  vec3 objectNormal = normalize(cross(tangent, bitangent));

  vNormal = objectNormal;
  vWorldPosition = modelPosition.xyz;

  gl_Position = projectionMatrix * viewMatrix * modelPosition;
}`

interface WaterOptions {
  environmentMap?: THREE.Texture;
  resolution?: number;
  size?: number;
  amplitude?:  number;
  frequency?:  number;
  persistence?:number;
  lacunarity?: number;
  iterations?: number;
  fresnelScale?: number;
  fresnelPower?: number;
}

export class Water extends THREE.Mesh{
  declare material: THREE.ShaderMaterial;
  constructor(options: WaterOptions = {}) {
    const geometry = new THREE.PlaneGeometry(
      options.size ?? 2,
      options.size ?? 2,
      options.resolution ?? 512,
      options.resolution ?? 512
    );

    const material = new THREE.ShaderMaterial({
      vertexShader: waterVert,
      fragmentShader: waterFrag,
      uniforms: {
        uTime: { value: 0 },
        uOpacity: { value: 0.8 },
        uEnvironmentMap: { value: options.environmentMap },
        uWavesAmplitude: { value: options.amplitude ?? 0.03 },
        uWavesFrequency: { value: options.frequency ?? 2.0 },
        uWavesPersistence: { value: options.persistence ?? 0.5 },
        uWavesLacunarity: { value: options.lacunarity ?? 2.0 },
        uWavesIterations: { value:  options.iterations ?? 5 },
        uWavesSpeed: { value: new THREE.Vector2(0.4, 0.4) },
        uTroughColor: { value: new THREE.Color('#186691') },
        uSurfaceColor: { value: new THREE.Color('#9bd8c0') },
        uPeakColor: { value: new THREE.Color('#bbd8e0') },
        uPeakThreshold: { value: 0.08 },
        uPeakTransition: { value: 0.05 },
        uTroughThreshold: { value: -0.01 },
        uTroughTransition: { value: 0.15 },
        uFresnelScale: { value: options.fresnelScale ?? 0.1 },
        uFresnelPower: { value:  options.fresnelPower ?? 2.0}
      },
      transparent: true,
      depthTest: true,
      side: THREE.DoubleSide
    });

    super(geometry, material);

    this.rotation.x = Math.PI * 0.5;
    this.position.y = 0;
  }

  update(time: number): void {
    this.material.uniforms.uTime.value = time;
  }

  initGUI(gui: GUI): void {
      const wavesFolder = gui.addFolder( 'Waves' );
      const uniforms = this.material.uniforms;
      const config = {
    uTroughColor: '#' + uniforms.uTroughColor.value.getHexString(),
    uSurfaceColor: '#' + uniforms.uSurfaceColor.value.getHexString(),
    uPeakColor: '#' + uniforms.uPeakColor.value.getHexString()
  }
  const colorFolder = gui.addFolder('Color');


  colorFolder.add(uniforms.uOpacity, 'value', 0, 1, 0.01,);
  colorFolder.addColor(config, 'uPeakColor').onChange((v) => {
    uniforms.uPeakColor.value.set(v)
  });
  colorFolder.addColor(config, 'uTroughColor').onChange((v) => {
    uniforms.uTroughColor.value.set(v)
  });
  colorFolder.addColor(config, 'uSurfaceColor').onChange((v) => {
    uniforms.uSurfaceColor.value.set(v)
  });


  colorFolder.add(uniforms.uPeakThreshold, 'value',
    0,
    0.5,
  ).name('Peak Threshold');

  colorFolder.add(uniforms.uPeakTransition, 'value', 0,
    0.5,
  ).name('Peak Transition');

  colorFolder.add(uniforms.uTroughThreshold, 'value', -0.5,
    0,
  ).name('Trough Threshold');

  colorFolder.add(uniforms.uTroughTransition, 'value', 0,
    0.5,
  ).name('Trough Transition');


  // Fresnel
  const fresnelFolder = gui.addFolder('Fresnel');
  fresnelFolder.add(uniforms.uFresnelScale, 'value', 0,
    1,
  ).name('Scale');

  fresnelFolder.add(uniforms.uFresnelPower, 'value', 0,
    3,
  ).name('Power');

  wavesFolder.add(uniforms.uWavesAmplitude, 'value',0, 1,).name( 'Amplitude')
  wavesFolder.add(uniforms.uWavesFrequency, 'value',0.1, 10,).name( 'frequency')
  wavesFolder.add(uniforms.uWavesPersistence, 'value',0, 1,).name( 'persistence')
  wavesFolder.add(uniforms.uWavesLacunarity, 'value',0, 3,).name( 'lacunarity')
  wavesFolder.add(uniforms.uWavesIterations, 'value',0, 10,).name( 'iterations')
  wavesFolder.add(uniforms.uWavesSpeed.value, 'x',-1, 1,).name( 'uWavesSpeedX')
  wavesFolder.add(uniforms.uWavesSpeed.value, 'y',-1, 1,).name( 'uWavesSpeedY')

  }
}
